"""
动动手吧：
1.卷积层增加regularizer，观察最终测试结果有什么变化
2.最终测试集accuracy 99.1%，实际上还没到lenet的最优点，能否通过调整超参数（lr、epoch、batchsize等）、优化器等使得优化结果更佳
3.能否修改网络结构，让mnist识别率更高
4.尝试使用GPU训练，与CPU比较
"""
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import numpy as np

# 载入mnist数据，如果指定目录下没有已经下载好的数据，则tensorflow会自动下载
mnist = input_data.read_data_sets("../mnist")
print("Training data size:", mnist.train.num_examples)
print("Validation data size:", mnist.validation.num_examples)
print("Testing data size:", mnist.test.num_examples)
BATCH_SIZE = 128
IMAGE_HEIGHT = 28
IMAGE_WIDTH = 28
IMAGE_CHANNELS = 1

# 网络结构
INPUT_NODE = 784
OUTPUT_NODE = 10
CONV1_DEEP = 32  # 第一卷积层的尺度和深度
CONV1_SIZE = 5  # 卷积核尺寸
CONV2_DEEP = 64  # 第二层卷积层的尺度和深度
CONV2_SIZE = 5  # 卷积核尺寸
FC_SIZE = 512
LEARNING_RATE_BASE = 0.1  # 初始化learning rate
LEARNING_RATE_DECAY = 0.99  # lr 衰减decay
EPOCH = 20  # 训练epoch迭代次数
TRAINING_STEPS = int(EPOCH * mnist.train.num_examples / BATCH_SIZE)  # 训练步数
configproto = tf.ConfigProto(allow_soft_placement=True)

# 定义了网络结构（静态图）
def inference(input_tensor, train, regularizer):
    """
    定义网络结构静态图
    :param input_tensor: 输入
    :param train: 是否是训练阶段
    :param regularizer: 正则
    :return: output
    """
    with tf.variable_scope('layer1-conv1', reuse=tf.AUTO_REUSE):
        conv1_weights = tf.get_variable("weight", [CONV1_SIZE, CONV1_SIZE, IMAGE_CHANNELS, CONV1_DEEP],
                                        initializer=tf.truncated_normal_initializer(stddev=0.1))
        conv1_biases = tf.get_variable("bias", [CONV1_DEEP], initializer=tf.constant_initializer(0.0))
        conv1 = tf.nn.conv2d(input_tensor, conv1_weights, strides=[1, 1, 1, 1], padding='SAME')
        relu1 = tf.nn.relu(tf.nn.bias_add(conv1, conv1_biases))
    with tf.name_scope('layer2-pool1'):
        pool1 = tf.nn.max_pool(relu1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    with tf.variable_scope('layer3-conv2', reuse=tf.AUTO_REUSE):
        conv2_weights = tf.get_variable('weight', [CONV2_SIZE, CONV2_SIZE, CONV1_DEEP, CONV2_DEEP],
                                        initializer=tf.truncated_normal_initializer(stddev=0.1))
        conv2_biases = tf.get_variable('bias', [CONV2_DEEP], initializer=tf.constant_initializer(0.0))
        conv2 = tf.nn.conv2d(pool1, conv2_weights, strides=[1, 1, 1, 1], padding='SAME')
        relu2 = tf.nn.relu(tf.nn.bias_add(conv2, conv2_biases))
    with tf.name_scope('layer4-pool2'):
        pool2 = tf.nn.max_pool(relu2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    pool_shape = pool2.get_shape().as_list()
    nodes = pool_shape[1] * pool_shape[2] * pool_shape[3]
    reshaped = tf.reshape(pool2, [pool_shape[0], nodes])
    with tf.variable_scope('layer5-fc1', reuse=tf.AUTO_REUSE):
        fc1_weights = tf.get_variable('weight', [nodes, FC_SIZE],
                                      initializer=tf.truncated_normal_initializer(stddev=0.1))
        if regularizer != None:
            tf.add_to_collection('losses', regularizer(fc1_weights))
        fc1_biases = tf.get_variable('bias', [FC_SIZE], initializer=tf.constant_initializer(0.1))
        fc1 = tf.nn.relu(tf.matmul(reshaped, fc1_weights) + fc1_biases)
        if train:
            fc1 = tf.nn.dropout(fc1, 0.5)
    with tf.variable_scope('layer6-fc2', reuse=tf.AUTO_REUSE):
        fc2_weights = tf.get_variable('weight', [FC_SIZE, OUTPUT_NODE],
                                      initializer=tf.truncated_normal_initializer(stddev=0.1))
        if regularizer != None:
            tf.add_to_collection('losses', regularizer(fc2_weights))
        fc2_biases = tf.get_variable('bias', [OUTPUT_NODE], initializer=tf.constant_initializer(0.1))
        logit = tf.matmul(fc1, fc2_weights) + fc2_biases
    return logit


def train(mnist):
    # input place holder
    x_input = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS], name='input')
    y_label = tf.placeholder(dtype=tf.float32, shape=[BATCH_SIZE, OUTPUT_NODE], name="label")

    # 建立计算图
    train_y = inference(x_input, True, tf.contrib.layers.l2_regularizer(0.0001))
    inference_y = inference(x_input, False, tf.contrib.layers.l2_regularizer(0.0001))
    prob = tf.nn.softmax(inference_y, name="inference-output")  # inference 的结果

    # 求取loss
    batch_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=train_y, labels=tf.argmax(y_label, axis=1))
    loss = tf.reduce_mean(batch_loss)
    loss = loss + tf.add_n(tf.get_collection(key="losses"))  # 把加入"loss"集合中的正则拿出来加入loss

    # 定义优化器
    global_step = tf.Variable(initial_value=0, trainable=False)
    learning_rate = tf.train.exponential_decay(LEARNING_RATE_BASE, global_step, mnist.train.num_examples / BATCH_SIZE,
                                               LEARNING_RATE_DECAY)
    optim = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss, global_step)  # 定义优化器

    with tf.Session(config=configproto) as sess:
        with tf.device("gpu:0"):
            sess.run(tf.global_variables_initializer())  # 初始化所有权重

            for step in range(TRAINING_STEPS):

                # 验证
                if step % 1000 == 0:
                    valid_steps = int(mnist.validation.num_examples / BATCH_SIZE)
                    whole_accuracy = []
                    for valid_step in range(valid_steps):
                        # 获取验证集
                        valid_input, valid_label= mnist.validation.next_batch(BATCH_SIZE)
                        valid_input = np.reshape(valid_input, (BATCH_SIZE, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS))
                        result = sess.run(prob, feed_dict={x_input: valid_input})
                        correct_prediction = np.equal(np.argmax(result, axis=1), np.argmax(valid_label, axis=1))
                        accuracy = np.mean(correct_prediction)
                        whole_accuracy.append(accuracy)
                    whole_accuracy = np.array(whole_accuracy).mean(axis=0)
                    print("After %d training step(s), validation accuracy using model is %g" % (step, whole_accuracy))

                # 训练
                train_data, labels = mnist.train.next_batch(BATCH_SIZE)  # 获取训练数据
                reshaped_xs = np.reshape(train_data, (BATCH_SIZE, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS))
                _, loss_ = sess.run([optim, loss], feed_dict={x_input: reshaped_xs, y_label: labels})
                print("After %d training step(s), training loss  is %g" % (step, loss_))

            # 测试
            test_steps = int(mnist.test.num_examples / BATCH_SIZE)
            whole_accuracy = []
            for test_step in range(test_steps):
                # 获取测试集
                test_input, test_label = mnist.test.next_batch(BATCH_SIZE)
                test_input = np.reshape(test_input, (BATCH_SIZE, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS))
                result = sess.run(prob, feed_dict={x_input: test_input})
                correct_prediction = np.equal(np.argmax(result, axis=1), np.argmax(test_label, axis=1))
                accuracy = np.mean(correct_prediction)
                whole_accuracy.append(accuracy)
            whole_accuracy = np.array(whole_accuracy).mean(axis=0)
            print("test accuracy using model is %g" % whole_accuracy)


if __name__ == "__main__":
    mnist = input_data.read_data_sets("../mnist", one_hot=True)  # 这里是tensorflow提供给mnist数据的专用类，如果是其他自己的数据
    train(mnist)
